package algorithm;

import algorithm.bean.Ant;

import java.util.ArrayList;
import java.util.List;

/**
 * 蚁群优化算法
 * @author letere
 * @since  2020-12-25 9:30
 */
public class AcoAlg {
    /**
     * 初始信息素大小
     */
    private static final double PHEROMONE_INIT_SIZE = 0.1;
    /**
     * 信息素挥发率
     */
    private static final double VOLATILIZATION_RATE = 0.5;

    /**
     * 距离矩阵
     */
    private final int[][] distanceMatrix;
    /**
     * 运行次数
     */
    private final int runTimes;
    /**
     * 蚂蚁数量
     */
    private final int antSize;

    /**
     * 信息素矩阵
     */
    private final double[][] pheromoneMatrix;
    /**
     * 最短路线
     */
    private List<Integer> minRoute;
    /**
     * 最短距离
     */
    private Integer minDistance = -1;

    public AcoAlg(int[][] distanceMatrix, int runTimes, int antSize) {
        this.distanceMatrix = distanceMatrix;
        this.runTimes = runTimes;
        this.antSize = antSize;
        // 初始化信息素矩阵
        int citySize = distanceMatrix.length;
        pheromoneMatrix = new double[citySize][citySize];
        for (int i=0; i<citySize; i++) {
            for (int j=0; j<citySize; j++) {
                pheromoneMatrix[i][j] = PHEROMONE_INIT_SIZE;
            }
        }
    }

    /**
     * 执行程序
     */
    public void execute() {
        this.runAlg();
        this.printResult();
    }

    /**
     * 运行蚁群优化算法
     */
    private void runAlg() {
        Ant ant;
        for(int i=0; i<runTimes; i++){
            List<Ant> ants = new ArrayList<>();
            for(int j=0; j<antSize; j++){
                ant = new Ant(pheromoneMatrix, distanceMatrix);
                ant.run();
                ants.add(ant);
            }
            // 更新信息素
            updatePheromone(ants);
            // 记录最佳路线与距离
            recordBest(ants);
        }
    }

    /**
     * 更新信息素
     * @param ants 蚂蚁群
     */
    private void updatePheromone(List<Ant> ants) {
        //信息素挥发
        for(int i=0; i<distanceMatrix.length; i++) {
            for(int j=0; j<distanceMatrix.length; j++) {
                pheromoneMatrix[i][j] = pheromoneMatrix[i][j] * (1.0 - VOLATILIZATION_RATE);
            }
        }
        //信息素新增
        for(Ant ant : ants) {
            List<Integer> route = ant.getRoute();
            for(int i=0; i<route.size()-1; i++){
                pheromoneMatrix[route.get(i)][route.get(i+1)] += 1.0 / ant.getRouteLength();
            }
        }
    }

    /**
     * 记录最短路径和距离
     */
    private void recordBest(List<Ant> ants) {
        // 初始化
        if(minDistance == -1.0) {
            minDistance = ants.get(0).getRouteLength();
            minRoute = ants.get(0).getRoute();
        }
        //遍历比较最短
        for(Ant ant : ants) {
            if(minDistance > ant.getRouteLength()) {
                minDistance = ant.getRouteLength();
                minRoute = ant.getRoute();
            }
        }
    }

    /**
     * 输出结果
     */
    private void printResult() {
        System.out.printf("蚂蚁数：%d, 运行次数：%d\n", antSize, runTimes);
        System.out.printf("最短距离为：%d\n", minDistance);
        this.resortRoute();
        System.out.printf("路线为：%s\n", minRoute.toString());
    }

    /**
     * 路线重新排序（为了方便观察，闭环都是以城市0为起始点）
     */
    private void resortRoute() {
        int zeroIndex = 0;
        for (Integer city : minRoute) {
            if (city == 0) {
                break;
            }
            zeroIndex ++;
        }
        if (zeroIndex == 0) {
            return;
        }
        // 重排
        List<Integer> newRoute = new ArrayList<>(minRoute.size());
        for (int i=zeroIndex; i<minRoute.size(); i++) {
            newRoute.add(minRoute.get(i));
        }
        for (int i=1; i<zeroIndex; i++) {
            newRoute.add(minRoute.get(i));
        }
        newRoute.add(minRoute.get(zeroIndex));
        // 替换路线
        minRoute = newRoute;
    }
}
